import random
import json
import csv
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import numpy as np
from typing import List, Tuple, Dict

class MazeGenerator:
    def __init__(self, size: int):
        """
        初始化迷宫生成器
        :param size: 迷宫尺寸，必须为奇数且 >= 7
        """
        if size < 7 or size % 2 == 0:
            raise ValueError("迷宫尺寸必须为奇数且不小于7")
        
        self.size = size
        self.maze = [['#' for _ in range(size)] for _ in range(size)]
        
        # 元素类型及其颜色映射
        self.element_colors = {
            '#': 'black',    # 墙壁
            ' ': 'white',    # 通路
            'S': 'green',    # 起点
            'E': 'red',      # 终点
            'G': 'gold',     # 资源(金币)
            'T': 'blue',     # 陷阱
            'L': 'orange',   # 机关
            'B': 'purple'    # BOSS
        }
    
    def generate_maze(self) -> List[List[str]]:
        """
        使用分治法生成迷宫
        """
        # 1. 初始化全地图为墙 (已在__init__中完成)
        
        # 2. 开始分治递归
        self._divide_and_conquer(0, 0, self.size, self.size, domain=1)
        
        # 3. 添加特殊元素
        self._add_special_elements()
        
        return self.maze
    
    def _divide_and_conquer(self, x: int, y: int, width: int, height: int, domain: int):
        """
        分治递归函数
        :param x: 区域左上角x坐标
        :param y: 区域左上角y坐标
        :param width: 区域宽度
        :param height: 区域高度
        """
        # # 创建交互式可视化窗口
        # plt.ion()
        # fig, ax = plt.subplots(figsize=(8, 8))
        
        # def update_visualization(title=""):
        #     ax.clear()
        #     # 绘制迷宫当前状态
        #     for i in range(self.size):
        #         for j in range(self.size):
        #             color = self.element_colors[self.maze[i][j]]
        #             ax.add_patch(patches.Rectangle((j, self.size-i-1), 1, 1, facecolor=color))
            
        #     ax.set_xlim(-0.5, self.size)
        #     ax.set_ylim(-0.5, self.size)
        #     ax.set_title(f"迷宫生成过程 - {title}")
        #     ax.grid(True)
        #     plt.draw()
        #     plt.pause(0.1)  # 暂停以便观察
        #     input("按回车键继续...")  # 等待用户输入
        
        # 3. 分治到3*3或4*4大小的时候处理
        if width <= 4 or height <= 4:
            if width == 3 and height == 3:
                # 在3x3区域中心创建通路
                center_x = x + width // 2
                center_y = y + height // 2
                self.maze[center_y][center_x] = ' '
                # update_visualization(f"处理3x3区域 ({x},{y})")
            elif width == 4 and height == 4:
                # 在4x4区域中心2x2区域随机选择3个点作为通路
                center_points = [(x+2,y+2), (x+1,y+2), (x+2,y+1), (x+1,y+1)]
                candidate_points = [center_points[domain-1]]
                center_points.remove(center_points[domain-1])
                random.shuffle(center_points)
                for cx, cy in candidate_points:
                    self.maze[cy][cx] = ' '
                for cx, cy in center_points[:2]:
                    self.maze[cy][cx] = ' '
                # update_visualization(f"处理4x4区域 ({x},{y})")
            return
        
        # 2. 定位好地图中心点，以四个象限做分治
        left_center_x = x + (width - 1) // 2
        left_center_y = y + (height - 1) // 2
        nxt_width = width // 2 + 1
        nxt_height = height // 2 + 1
        
        if(width % 2 == 0):
            right_center_x = left_center_x + 1
            right_center_y = left_center_y + 1
            nxt_width = width // 2
            nxt_height = height // 2
        else:
            right_center_x = left_center_x
            right_center_y = left_center_y
        
        
        
        # update_visualization(f"开始处理区域 ({x},{y}) - 大小: {width}x{height}")
        
        # 递归处理四个象限
        # 左上象限
        self._divide_and_conquer(x, y, nxt_width, nxt_height, domain=1)
        # 右上象限
        self._divide_and_conquer(right_center_x, y, nxt_width, nxt_height, domain=2)
        # 左下象限
        self._divide_and_conquer(x, right_center_y, nxt_width, nxt_height, domain=3)
        # 右下象限
        self._divide_and_conquer(right_center_x, right_center_y, nxt_width, nxt_height, domain=4)
        
        # 4. 当这一层递归结束后，创建通道
        self._create_passages(left_center_x, left_center_y, x, y, width, height)
        # update_visualization(f"完成区域 ({x},{y}) 的通道创建")
        
        # plt.close()  # 关闭当前图形窗口
    
    def _create_passages(self, center_x: int, center_y: int, 
                        region_x: int, region_y: int, width: int, height: int):
        """
        在十字分割线上创建通道，选择3个方向开口，确保通道两侧都有通路
        """
        import time
        random.seed(int(time.time()) + random.randint(0, 1000000))
        directions = ['up', 'down', 'left', 'right']
        random.shuffle(directions)
        
        # # 选择3个方向开口
        # chosen_directions = directions[:3]
        
        valid = []
        
        
        offset = 1 if width % 2 == 1 else 2
        for direction in directions:
            if direction == 'up':
                # 在上方开口，确保开口两侧有通路
                valid_passages = []
                for x in range(region_x+1, center_x):
                    # 检查开口上下是否都有通路
                    if (self.maze[center_y-1][x] == ' ' and 
                        self.maze[center_y+offset][x] == ' '):
                        valid_passages.append(x)
                
                if valid_passages:
                    random.shuffle(valid_passages)
                    passage_x = valid_passages[0]
                    if 0 <= passage_x < self.size and 0 <= center_y < self.size:
                        valid.append((center_y, passage_x))
                        if offset == 2:
                            valid.append((center_y+1, passage_x))
            
            elif direction == 'down':
                # 在下方开口，确保开口两侧有通路
                valid_passages = []
                for x in range(center_x + 1, region_x + width-1):
                    if (self.maze[center_y-1][x] == ' ' and 
                        self.maze[center_y+offset][x] == ' '):
                        valid_passages.append(x)
                
                if valid_passages:
                    random.shuffle(valid_passages)
                    passage_x = valid_passages[0]
                    if 0 <= passage_x < self.size and 0 <= center_y < self.size:
                        valid.append((center_y, passage_x))
                        if offset == 2:
                            valid.append((center_y+1, passage_x))
            
            elif direction == 'left':
                # 在左侧开口，确保开口两侧有通路
                valid_passages = []
                for y in range(region_y+1, center_y):
                    if (self.maze[y][center_x-1] == ' ' and 
                        self.maze[y][center_x+offset] == ' '):
                        valid_passages.append(y)
                
                if valid_passages:
                    random.shuffle(valid_passages)
                    passage_y = valid_passages[0]
                    if 0 <= center_x < self.size and 0 <= passage_y < self.size:
                        valid.append((passage_y, center_x))
                        if offset == 2:
                            valid.append((passage_y, center_x+1))
                            
            elif direction == 'right':
                # 在右侧开口，确保开口两侧有通路
                valid_passages = []
                for y in range(center_y + 1, region_y + height - 1):
                    if (self.maze[y][center_x-1] == ' ' and 
                        self.maze[y][center_x+offset] == ' '):
                        valid_passages.append(y)
                
                if valid_passages:
                    random.shuffle(valid_passages)
                    passage_y = valid_passages[0]
                    if 0 <= center_x < self.size and 0 <= passage_y < self.size:
                        valid.append((passage_y, center_x))
                        if offset == 2:
                            valid.append((passage_y, center_x+1))
        
        print(valid)
        for y, x in valid[:offset*3]:
            print(f"创建通道: ({y}, {x})")
            self.maze[y][x] = ' '
    
    def _add_special_elements(self):
        """
        添加特殊元素：起点、终点、资源、陷阱、机关、BOSS
        """
        # 获取所有通路位置
        passages = []
        for i in range(self.size):
            for j in range(self.size):
                if self.maze[i][j] == ' ':
                    passages.append((i, j))
        
        if len(passages) < 6:  # 确保有足够的通路放置元素
            return
        
        random.shuffle(passages)
        
        # 放置必需元素
        elements_to_place = ['S', 'E']  # 起点和终点是必需的
        
        # 根据迷宫大小决定其他元素数量
        if self.size >= 15:
            elements_to_place.extend(['G'] * 4 + ['T'] * 4 + ['L'] * 1 + ['B'] * 1)
        elif self.size >= 11:
            elements_to_place.extend(['G'] * 3 + ['T'] * 3 + ['L'] * 1 + ['B'] * 1)
        else:  # 7x7 或更小
            elements_to_place.extend(['G'] * 2 + ['T'] * 2 + ['L'] * 1)
        
        # 确保不超过可用通路数量
        elements_to_place = elements_to_place[:len(passages)]
        
        # 随机分布元素
        for i, element in enumerate(elements_to_place):
            row, col = passages[i]
            self.maze[row][col] = element
    
    def save_to_json(self, filename: str = "./maze.json"):
        """保存迷宫到JSON文件"""
        with open(filename, 'w', encoding='utf-8') as f:
            json.dump({"maze": self.maze}, f, indent=2, ensure_ascii=False)
        print(f"迷宫已保存到 {filename}")
    
    def save_to_csv(self, filename: str = "./maze.csv"):
        """保存迷宫到CSV文件"""
        with open(filename, 'w', newline='', encoding='utf-8') as f:
            writer = csv.writer(f)
            writer.writerows(self.maze)
        print(f"迷宫已保存到 {filename}")
    
    def visualize(self, save_path: str = "./maze_visualization.png"):
        """
        可视化迷宫
        """
        fig, ax = plt.subplots(1, 1, figsize=(10, 10))
        
        # 绘制迷宫
        for i in range(self.size):
            for j in range(self.size):
                element = self.maze[i][j]
                color = self.element_colors.get(element, 'gray')
                
                # 创建方块
                rect = patches.Rectangle((j, self.size-1-i), 1, 1, 
                                       linewidth=1, edgecolor='gray', 
                                       facecolor=color)
                ax.add_patch(rect)
                
                # 在方块中添加文字标签
                if element != ' ' and element != '#':
                    ax.text(j + 0.5, self.size-1-i + 0.5, element, 
                           ha='center', va='center', fontsize=12, 
                           fontweight='bold', color='white' if color in ['black', 'blue', 'purple'] else 'black')
        
        # 设置坐标轴
        ax.set_xlim(0, self.size)
        ax.set_ylim(0, self.size)
        ax.set_aspect('equal')
        ax.set_title(f'Maze Visualization ({self.size}x{self.size})', fontsize=16, fontweight='bold')
        
        # 移除坐标轴刻度
        ax.set_xticks([])
        ax.set_yticks([])
        
        # 添加图例
        legend_elements = []
        for element, color in self.element_colors.items():
            if element == ' ':
                continue
            label = {'#': 'Wall', 'S': 'Start', 'E': 'Exit', 'G': 'Gold', 
                    'T': 'Trap', 'L': 'Locker', 'B': 'BOSS'}.get(element, element)
            legend_elements.append(patches.Patch(color=color, label=label))
        
        ax.legend(handles=legend_elements, loc='center left', bbox_to_anchor=(1, 0.5))
        
        plt.tight_layout()
        plt.savefig(save_path, dpi=300, bbox_inches='tight')
        plt.show()
        print(f"可视化图像已保存到 {save_path}")
    
    def print_maze(self):
        """打印迷宫到控制台"""
        print("\n生成的迷宫:")
        print("=" * (self.size * 2 + 1))
        for row in self.maze:
            print("|" + "|".join(row) + "|")
        print("=" * (self.size * 2 + 1))


def main():
    """主函数，演示迷宫生成"""
    print("=== 分治法迷宫生成器 ===")
    
    # 获取用户输入
    while True:
        try:
            size = int(input("请输入迷宫尺寸 (必须为奇数且 >= 7): "))
            if size < 7 or size % 2 == 0:
                print("错误：迷宫尺寸必须为奇数且不小于7")
                continue
            break
        except ValueError:
            print("请输入有效的数字")
    
    # 生成迷宫
    print(f"\n正在生成 {size}x{size} 迷宫...")
    generator = MazeGenerator(size)
    maze = generator.generate_maze()
    
    # 显示结果
    generator.print_maze()
    
    # 保存文件
    generator.save_to_json()
    generator.save_to_csv()
    
    # 可视化
    print("\n正在生成可视化图像...")
    generator.visualize()
    
    print("\n迷宫生成完成！")


if __name__ == "__main__":
    main() 